/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.ignite.internal.benchmarks.jmh.direct;

import java.nio.ByteBuffer;
import org.apache.ignite.internal.benchmarks.jmh.JmhAbstractBenchmark;
import org.apache.ignite.internal.benchmarks.jmh.runner.JmhIdeBenchmarkRunner;
import org.apache.ignite.internal.direct.DirectMessageWriter;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.profile.GCProfiler;
import org.openjdk.jmh.runner.options.TimeValue;

/**
 * Benchmarks for {@link DirectMessageWriter} performance changes:
 * <ul>
 *     <li>{@link #hotPathPrimitiveFields} measures the per-field write overhead (the {@code curStream} caching that
 *     avoids re-resolving {@code state.item().stream} on every primitive write). Run it on the baseline build and on
 *     the patched build and compare throughput.</li>
 *     <li>{@link #compressedScratchOld} vs {@link #compressedScratchNew} model the per-field scratch allocation done by
 *     {@code writeCompressedMessage}: a fresh direct buffer + a brand-new writer (old) versus a reused writer + heap
 *     buffer (new). Run together with {@code -prof gc} to see the allocation-rate / direct-memory difference.</li>
 * </ul>
 */
@State(Scope.Thread)
public class JmhDirectMessageWriterBenchmark extends JmhAbstractBenchmark {
    /** Network buffer the hot-path benchmark serializes into (reused, like the real comms write loop). */
    private static final int NET_BUF = 64 * 1024;

    /** Scratch buffer capacity, matches {@code DirectMessageWriter.TMP_BUF_CAPACITY}. */
    private static final int TMP_BUF_CAPACITY = 10 * 1024;

    /** Writer under test for the hot path. */
    private DirectMessageWriter writer;

    /** Reused network buffer. */
    private ByteBuffer buf;

    /** Reusable writer modeling the patched compressed path. */
    private DirectMessageWriter reusable;

    /** */
    @Setup
    public void setup() {
        // The message factory is only used by writeMessage(); primitive writes and scratch allocation never touch it.
        writer = new DirectMessageWriter(null);
        buf = ByteBuffer.allocate(NET_BUF);
        reusable = new DirectMessageWriter(null);
    }

    /**
     * Hot path: serialize a batch of mixed primitive fields into a reused buffer, mirroring how a generated serializer
     * calls the writer field by field. Stresses the {@code state.item().stream} resolution done on every call.
     *
     * @return Accumulated {@code lastFinished} flag (consumed to prevent dead-code elimination).
     */
    @Benchmark
    public boolean hotPathPrimitiveFields() {
        ByteBuffer b = buf;

        b.clear();

        writer.setBuffer(b);

        boolean ok = true;

        // 256 "messages" of 7 mixed fields = 1792 writer calls, ~7 KB, fits the 64 KB buffer.
        for (int i = 0; i < 256; i++) {
            ok &= writer.writeByte((byte)i);
            ok &= writer.writeShort((short)i);
            ok &= writer.writeInt(i);
            ok &= writer.writeLong(i);
            ok &= writer.writeBoolean((i & 1) == 0);
            ok &= writer.writeFloat(i);
            ok &= writer.writeDouble(i);
        }

        return ok;
    }

    /**
     * Old compressed-path scratch allocation: fresh direct buffer + brand-new writer per compressed field.
     *
     * @return Allocated writer (returned to prevent dead-code elimination).
     */
    @Benchmark
    public Object compressedScratchOld() {
        ByteBuffer tmp = ByteBuffer.allocateDirect(TMP_BUF_CAPACITY);

        DirectMessageWriter w = new DirectMessageWriter(null);

        w.setBuffer(tmp);

        return w;
    }

    /**
     * New compressed-path scratch allocation: reused writer + heap buffer per compressed field.
     *
     * @return Reused writer (returned to prevent dead-code elimination).
     */
    @Benchmark
    public Object compressedScratchNew() {
        ByteBuffer tmp = ByteBuffer.allocate(TMP_BUF_CAPACITY);

        DirectMessageWriter w = reusable;

        w.reset();
        w.setBuffer(tmp);

        return w;
    }

    /**
     * @param args Arguments.
     * @throws Exception If failed.
     */
    public static void main(String[] args) throws Exception {
        JmhIdeBenchmarkRunner.create()
            .forks(1)
            .threads(1)
            .warmupIterations(6)
            .warmupTime(TimeValue.milliseconds(500))
            .measurementIterations(10)
            .measurementTime(TimeValue.milliseconds(500))
            .benchmarks(JmhDirectMessageWriterBenchmark.class.getSimpleName())
            .jvmArguments("-Xms2g", "-Xmx2g")
            .profilers(GCProfiler.class)
            .run();
    }
}